import os

from dotenv import load_dotenv
from langchain_core.messages import AIMessage
from langchain_ollama import ChatOllama
import requests
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langgraph.constants import END, START
from langgraph.graph import MessagesState, StateGraph
from langgraph.prebuilt import create_react_agent, ToolNode
from typing import Literal


def download_db():
    global url, response, file
    url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"
    response = requests.get(url)
    if response.status_code == 200:
        # Open a local file in binary write mode
        with open("Chinook.db", "wb") as file:
            # Write the content of the response (the file) to the local file
            file.write(response.content)
        print("File downloaded and saved as Chinook.db")
    else:
        print(f"Failed to download the file. Status code: {response.status_code}")


def list_tables(state: MessagesState):
    tool_call = {
        "name": "sql_db_list_tables",
        "args": {},
        "id": "abc123",
        "type": "tool_call",
    }
    tool_call_message = AIMessage(content="", tool_calls=[tool_call])

    list_tables_tool = next(tool for tool in tools if tool.name == "sql_db_list_tables")
    tool_message = list_tables_tool.invoke(tool_call)
    response = AIMessage(f"Available tables: {tool_message.content}")

    return {"messages": [tool_call_message, tool_message, response]}

    # Example: force a model to create a tool call


def call_get_schema(state: MessagesState):
    # Note that LangChain enforces that all models accept `tool_choice="any"`
    # as well as `tool_choice=<string name of tool>`.
    llm_with_tools = model.bind_tools([get_schema_tool], tool_choice="any")
    response = llm_with_tools.invoke(state["messages"])

    return {"messages": [response]}


def generate_query(state: MessagesState):
    system_message = {
        "role": "system",
        "content": generate_query_system_prompt,
    }
    # We do not force a tool call here, to allow the model to
    # respond naturally when it obtains the solution.
    llm_with_tools = model.bind_tools([run_query_tool])
    response = llm_with_tools.invoke([system_message] + state["messages"])

    return {"messages": [response]}


def check_query(state: MessagesState):
    system_message = {
        "role": "system",
        "content": check_query_system_prompt,
    }

    # Generate an artificial user message to check
    tool_call = state["messages"][-1].tool_calls[0]
    user_message = {"role": "user", "content": tool_call["args"]["query"]}
    llm_with_tools = model.bind_tools([run_query_tool], tool_choice="any")
    response = llm_with_tools.invoke([system_message, user_message])
    response.id = state["messages"][-1].id

    return {"messages": [response]}


def should_continue(state: MessagesState) -> Literal[END, "check_query"]:
    messages = state["messages"]
    last_message = messages[-1]
    if not last_message.tool_calls:
        return END
    else:
        return "check_query"


if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    model = ChatOllama(model="qwen3:30b", base_url="http://192.168.97.217:11434")
    # download_db()
    db = SQLDatabase.from_uri("sqlite:///Chinook.db")
    print(f"Dialect: {db.dialect}")
    print(f"Available tables: {db.get_usable_table_names()}")
    print(f'Sample output: {db.run("SELECT * FROM Artist LIMIT 5;")}')

    toolkit = SQLDatabaseToolkit(db=db, llm=model)

    tools = toolkit.get_tools()
    for tool in tools:
        print(f"{tool.name}: {tool.description}\n")

    # sql_db_query
    # 该工具的输入是一条详细且正确的 SQL 查询语句，输出是来自数据库的查询结果。 如果 SQL 查询不正确，将返回错误信息。
    # 如果返回错误，请重写查询语句，检查语法，并重试。 如果出现 Unknown column 'xxxx' in 'field list' 的问题，请使用 sql_db_schema 工具查询正确的表字段。
    # sql_db_schema
    # 该工具的输入是用英文逗号分隔的表名列表，输出是这些表的表结构（schema）和部分示例行数据。
    # 请确保这些表实际存在，必要时先调用 sql_db_list_tables 工具！
    # 示例输入：table1, table2, table3
    # sql_db_list_tables
    # 该工具的输入为空字符串，输出是数据库中表名的英文逗号分隔列表。
    # sql_db_query_checker
    # 使用该工具可以在执行 SQL 查询前进行语法检查，确保查询正确。 在使用 sql_db_query 执行查询前，务必先调用此工具进行校验！
    system_prompt = """
        你是一个专为与 SQL 数据库交互而设计的智能代理。
        给定一个输入问题，请生成一个语法正确的 {dialect} 查询语句以执行，
        然后查看查询结果并返回答案。除非用户明确要求获取特定数量的结果，
        否则查询结果请始终限制在最多 {top_k} 条以内。

        你可以通过某个相关列对结果进行排序，以返回数据库中最有价值的结果。
        切勿查询某张表中的所有列，只查询与问题相关的列。

        在执行查询之前，你必须仔细检查你的 SQL 语句。如果执行时发生错误，
        请重新编写查询语句并再次尝试。

        不要对数据库执行任何 DML 语句（如 INSERT、UPDATE、DELETE、DROP 等）。

        一开始你必须查看数据库中的所有表，以了解可以查询哪些内容，
        绝不能跳过这一步。

        接下来你应该查询与问题最相关表的表结构（schema）。
        """.format(
        dialect=db.dialect,
        top_k=5,
    )

    agent = create_react_agent(
        model,
        tools,
        prompt=system_prompt,
    )

    question = "平均而言，哪种音乐类型的曲目最长?"

    for step in agent.stream(
            {"messages": [{"role": "user", "content": question}]},
            stream_mode="values",
    ):
        step["messages"][-1].pretty_print()

    get_schema_tool = next(tool for tool in tools if tool.name == "sql_db_schema")
    get_schema_node = ToolNode([get_schema_tool], name="get_schema")

    run_query_tool = next(tool for tool in tools if tool.name == "sql_db_query")
    run_query_node = ToolNode([run_query_tool], name="run_query")

    # Example: create a predetermined tool call

    generate_query_system_prompt = """
    你是一个专为与 SQL 数据库交互而设计的智能代理。
    根据输入的问题，生成一条语法正确的 {dialect} 查询语句用于执行，
    然后查看查询结果并返回答案。除非用户明确指定需要获取的示例数量，
    否则始终将查询结果限制为最多 {top_k} 条。

    你可以根据相关列对结果进行排序，以返回数据库中最有代表性的示例。
    切勿查询某张表的所有列，仅查询与问题相关的列。

    严禁对数据库执行任何 DML 操作（如 INSERT、UPDATE、DELETE、DROP 等）。
    """.format(
        dialect=db.dialect,
        top_k=5,
    )

    check_query_system_prompt = """
    你是一位对细节高度关注的 SQL 专家。
    请仔细检查该 {dialect} 查询语句是否存在常见错误，包括：

    - 使用 NOT IN 时未考虑 NULL 值的影响
    - 使用 UNION 而不是更合适的 UNION ALL
    - 在需要排除区间时使用了 BETWEEN
    - 条件中数据类型不匹配
    - 标识符是否正确加引号
    - 函数参数数量是否正确
    - 是否正确进行类型转换
    - 连接（JOIN）语句中是否使用了正确的列

    如果存在上述任何错误，请重写 SQL 查询；
    如果没有错误，请原样返回原始查询语句。

    你将在检查之后调用相应工具来执行该查询。
    """.format(dialect=db.dialect)

    builder = StateGraph(MessagesState)
    builder.add_node(list_tables)
    builder.add_node(call_get_schema)
    builder.add_node("get_schema", get_schema_node)
    builder.add_node(generate_query)
    builder.add_node(check_query)
    builder.add_node("run_query", run_query_node)

    builder.add_edge(START, "list_tables")
    builder.add_edge("list_tables", "call_get_schema")
    builder.add_edge("call_get_schema", "get_schema")
    builder.add_edge("get_schema", "generate_query")
    builder.add_conditional_edges(
        "generate_query",
        should_continue,
    )
    builder.add_edge("check_query", "run_query")
    builder.add_edge("run_query", "generate_query")

    agent = builder.compile()

    graph_image = agent.get_graph().draw_mermaid_png()
    with open("sql_agent.png", "wb") as f:
        f.write(graph_image)

    question = "平均而言，哪种音乐类型的曲目最长?"

    for step in agent.stream(
            {"messages": [{"role": "user", "content": question}]},
            stream_mode="values",
    ):
        step["messages"][-1].pretty_print()
